package com.sz.szaiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 查询重写器
 * 该类用于对用户输入的查询进行重写，以提升查询的准确性和语义理解。
 */
@Component
public class QueryRewriter {

    private final QueryTransformer queryTransformer;

    /**
     * 构造函数，初始化查询重写转换器
     * @param dashscopeChatModel 用于查询重写的AI模型
     */
    public QueryRewriter(ChatModel dashscopeChatModel) {
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel); // 创建ChatClient构建器
        // 创建查询重写转换器
        queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder) // 设置ChatClient构建器
                .build(); // 构建RewriteQueryTransformer实例
    }

    /**
     * 执行查询重写
     * @param prompt 用户输入的原始查询
     * @return 重写后的查询文本
     */
    public String doQueryRewrite(String prompt) {
        Query query = new Query(prompt); // 创建查询对象
        // 执行查询重写
        Query transformedQuery = queryTransformer.transform(query); // 应用查询重写转换器
        // 输出重写后的查询
        return transformedQuery.text(); // 返回重写后的查询文本
    }
}
